#!/usr/bin/env python
import numpy as np
from astropy.io import fits
from glob import glob
from scipy.stats import sigmaclip
from astropy.table import Table
from scipy.ndimage import median_filter
from argparse import ArgumentParser
from scipy import odr
import matplotlib.pyplot as plt
import os
from scipy.stats import sigmaclip

class e2v290:
    def __init__(self):
        self.n_channel = 16
        self.col = [0, 1, 2, 3, 4, 5, 6, 7, 7, 6, 5, 4, 3, 2, 1, 0]
        self.row = [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1]
        self.add_x = [1250 * i for i in self.col]
        self.add_y = [4700 * i for i in self.row]
        
class cmos:
    def __init__(self):
        # channel layout
        self.n_channel = 2
        self.col = [0, 1]
        self.row = [0, 0]
        self.add_x = [0, 4988]
        self.add_y = [0, 0]
        
def det_class(det_type):
    valid = ['e2v290','cmos']
    if det_type == 'e2v290':
        return e2v290()
    if det_type == 'cmos':
        return cmos()
        
def get_ptc_data(fdir,xc,yc,xs,ys,dc,rdir):
    bias_files = glob(fdir+"/ptc-0-*.fits")#
    ptc_files = []
    for i in range(60):
        if not os.path.exists(fdir+"/ptc-"+str(i+1)+"-1.fits"):#
            break
        ptc_files += glob(fdir+"/ptc-"+str(i+1)+"-*.fits")
    print(ptc_files)
    for i in range(len(bias_files)):
        if i==0:
            bias = fits.getdata(bias_files[i]).astype('float')
        else:
            bias += fits.getdata(bias_files[i]).astype('float')
    bias /= len(bias_files)
    l = int(len(ptc_files)/2)
    #l = 56
    ptc_data = np.zeros((l,dc.n_channel,2))
    rc_data = np.zeros((l,dc.n_channel))
    region_check = np.zeros(fits.getdata(fdir+"/ptc-"+str(i+1)+"-1.fits").shape)#
    for j in range(dc.n_channel):
        xcent = xc + dc.add_x[j]
        ycent = yc + dc.add_y[j]
        region_check[int(ycent-ys/2):int(ycent+ys/2),int(xcent-xs/2):int(xcent+xs/2)] = 1
    fits.writeto(rdir+"/region_check.fits",data = region_check.astype('uint8'),overwrite=True)
    print("l",l)
    for i in range(l-1):
        im1 = fits.getdata(fdir+"/ptc-"+str(i+1)+"-1.fits").astype('float')#
        im2 = fits.getdata(fdir+"/ptc-"+str(i+1)+"-2.fits").astype('float')#
        if i==0:
            wave = str(fits.open(fdir+"/ptc-"+str(i+1)+"-1.fits")[0].header['WAVELEN'])#
        for j in range(dc.n_channel):
            xcent = xc + dc.add_x[j]
            ycent = yc + dc.add_y[j]
            im1_cut = im1[int(ycent-ys/2):int(ycent+ys/2),int(xcent-xs/2):int(xcent+xs/2)]
            im2_cut = im2[int(ycent-ys/2):int(ycent+ys/2),int(xcent-xs/2):int(xcent+xs/2)]
            bias_cut = bias[int(ycent-ys/2):int(ycent+ys/2),int(xcent-xs/2):int(xcent+xs/2)]
            #print("im1",np.mean(im1_cut),"im2",np.mean(im2_cut),"bias",np.mean(bias_cut))
            dn_raw = (im1_cut+im2_cut)/2 - bias_cut
            ptc_data[i,j,0] = np.mean((im1_cut + im2_cut)/2 - bias_cut)
            ptc_data[i,j,1] = np.var(im1_cut-im2_cut)/2
            rc_data[i,j] = ptc_data[i,j,0]
    print(ptc_data.shape)
    return ptc_data,rc_data,wave

def plot_ptc(x,y,nonlin,fw,valid,gain,plot_dir,channel,idx20,idx10,wave):
    fig,axes = plt.subplots(2,1,figsize=(8,14))
    axes[0].scatter(x,y,c='orange',s=20)
    axes[0].scatter(x[valid],y[valid],c='g',s=20,label='used data')
    axes[0].plot(x,x/gain,label="fit",c='orange')
    axes[0].legend()
    #print(nonlin[1:])
    axes[1].scatter(x[1:],nonlin[1:],c='orange',s=20)
    axes[1].scatter(x[idx20],nonlin[idx20],c='b',s=20,label='14-56ke-')
    axes[1].scatter(x[idx10],nonlin[idx10],c='r',s=20,label='7-70ke-')
    axes[1].axhline(y=3, c='grey', linestyle='--')
    axes[1].axhline(y=-3, c='grey', linestyle='--')
    axes[1].axhline(y=5, c='grey', linestyle='--')
    axes[1].axhline(y=-5, c='grey', linestyle='--')
    axes[1].axvline(x=fw, c='blue', linestyle='--', label='full '+str(np.round(fw*gain,3))+"e-")
    axes[1].legend()
    fig.savefig(plot_dir+"/ptc_"+wave+"_c"+str(channel)+".jpg")
    fig.clf()
    
def fit_ptc(ptc_data,result_dir,wave):
    tab = Table()
    tab['channel'] = np.arange(ptc_data.shape[1])+1
    mod = odr.polynomial(1)
    fws = np.zeros(ptc_data.shape[1])
    nlptc10s = np.zeros(ptc_data.shape[1])
    nlptc20s = np.zeros(ptc_data.shape[1])
    gains = np.zeros(ptc_data.shape[1])
    #print(len(gains))
    for j in range(ptc_data.shape[1]):
        x = ptc_data[:,j,0]
        y = ptc_data[:,j,1]
        if ptc_data.shape[0] > 10:
            x = median_filter(x, 3)
            peak_x = np.argmax(x)
            peak_y = np.argmax(y)
            peak = min(peak_x, peak_y)
            fmax = x[peak]
            valid = (x<=fmax)
        else:
            peak_x = np.argmax(x)
            peak_y = np.argmax(y)
            peak = min(peak_x, peak_y)
            fmax = x[peak]
            valid = (x <= fmax)
        print("fmax",fmax,"xmax",np.max(x),"valid",valid)
        xfit = x[valid]
        yfit = y[valid]
        data = odr.Data(xfit, yfit, wd=1 / np.power(xfit, 2), we=1 / np.power(yfit, 2))
        fitter_ptc = odr.ODR(data, mod, ifixb=[0, 1], beta0=[0, np.median(yfit / xfit)])
        res = fitter_ptc.run()
        gains[j] = 1/res.beta[1]
        fullwell = 7e4
        fw_use = fullwell / gains[j]
        idx20 = (x < fw_use * 0.8) * (x > fw_use * 0.2)
        idx10 = (x < fw_use) * (x > fw_use * 0.1)
        
        fw = np.nan
        nonlin_ptc_10 = np.nan
        nonlin_ptc_20 = np.nan
        
        # full well
        nonlin = (y / (x / gains[j]) - 1) * 100
        index = (x > fmax * 0.3) * (x <= fmax) * (np.abs(nonlin) > 5)
        if index.sum() == 0:
            fw = fmax
        else:
            p2 = min(np.arange(len(x), dtype=int)[index])
            p1 = p2 - 1
            if nonlin[p2] > 5:
                fw = np.interp(5, [nonlin[p1], nonlin[p2]], [x[p1], x[p2]])
            else:
                fw = np.interp(-5, [nonlin[p1], nonlin[p2]], [x[p1], x[p2]])
        
        #plot PTC
        plot_ptc(x,y,nonlin,fw,valid,gains[j],result_dir,j+1,idx20,idx10,wave)
        
        # PTC non-linearity
        nonlin_ptc_10 = np.max(abs(nonlin[np.where(idx10*valid)]))
        print(idx10*valid)
        print("nlptc10",nonlin_ptc_10)
        nonlin_ptc_20 = np.max(abs(nonlin[np.where(idx20*valid)]))
        fws[j] = fw*gains[j]
        nlptc10s[j] = nonlin_ptc_10
        nlptc20s[j] = nonlin_ptc_20
    
    tab['gain'] =  gains
    tab['gain'].unit = 'e-/ADU'
    tab['fullwell'] = fws
    tab['fullwell'].unit = 'e-'
    tab['nlptc20'] = nlptc20s
    tab['nlptc20'].unit = '%'
    tab['nlptc10'] = nlptc10s
    tab['nlptc10'].unit = '%'
    return tab
    
def _run(fdir,result_dir,xc,yc,xs,ys,det_type):
    dc = det_class(det_type)
    ptc_data,rc_data,wave = get_ptc_data(fdir,xc,yc,xs,ys,dc,result_dir)
    #ptc_tab = fit_ptc(ptc_data,result_dir,wave)
    ptc_tab.write(result_dir+"/ptc_result_"+wave+"nm.tab",format='ipac',overwrite=True)
    return  

def main():
    parser = ArgumentParser()
    parser.add_argument('datadir', type=str)
    parser.add_argument('outdir', type=str)
    parser.add_argument('--xcenter', '-xc', type=int, default=None)
    parser.add_argument('--ycenter', '-yc', type=int, default=None)
    parser.add_argument('--xsize', '-xs', type=int, default=None)
    parser.add_argument('--ysize', '-ys', type=int, default=None)
    parser.add_argument('--detector_type', '-dt', default='e2v290')
    args = parser.parse_args()
    _run(args.datadir,args.outdir,args.xcenter,args.ycenter,args.xsize,args.ysize,args.detector_type)

if __name__ == '__main__':
    main()
